# coding=utf-8
import codecs
import json
import logging
import yaml
from typing import List, Dict, Hashable
from collections import OrderedDict
from trie import Trie
from transformers import AutoTokenizer

logger = logging.getLogger(__name__)


def build_tries(dict_config_file: str = 'dict_config.yaml') -> Dict[Hashable, Trie]:
    """
    依据 dict_config 词典配置文件读取所有词典，每种词典类别分别构建各自的字典树 Trie，最终返回所有的 Tries。

    Args:
        dict_config_file (yaml): 词典配置文件，用于描述词典存放路径
    Returns:
        Dict: {词典名：字典树}
    """
    tries = {}

    with codecs.open(dict_config_file, 'r', 'utf-8') as dict_config:
        config = yaml.load(dict_config, Loader=yaml.FullLoader)
        for dict_name in config:
            dict_file = codecs.open(config[dict_name], 'r', 'utf-8')
            dict_words = set()
            for line in dict_file: dict_words.add(line.strip())
            trie = Trie(words=list(dict_words))
            tries[dict_name] = trie

    return tries


def tries_search(tries, tokens):
    res = []

    for trie_name in tries:
        output = tries[trie_name].search(tokens)
        if output:
            res.extend([{trie_name: list(output.values())}])

    return res


def convert_lookup_to_cgr(lookup_res: List, tag_dict, sent_lens, other_marker=True):
    """
    参考 GemNet 论文，希望把离线词典信息变成论文中要求的 (Contextual Gazetteer Representation)CGR 格式。
    ！这里有个很重要的分支！
    对于不再此表的词典，到底是 tag=O 上为 1，还是全部为 0。
    第一种相当于让所有不在离线词典中的词语，都为黑名单，被标记为 O，这种是不是太强的信号？或许离线词典维护完善场景更适合这种设计？
    第二种，相当于自己要有个专门的黑名单，防止模型误学不应该是对应领域的词语信息。

    XX    | O | B-Per | I-Per | B-Citizen | I-Citizen | B-City | I-City |
    小    | 0 |   1   |   0   |     0     |     0     |    0   |   0    |
    明    | 0 |   0   |   1   |     0     |     0     |    0   |   0    |
    *是*  | 1 |   0   |   0   |     0     |     0     |    0   |   0    |
    中    | 0 |   0   |   0   |     1     |     0     |    1   |   0    |
    国    | 0 |   0   |   0   |     0     |     1     |    0   |   1    |
    人    | 0 |   0   |   0   |     0     |     1     |    0   |   0    |
    。    | 1 |   0   |   0   |     0     |     0     |    0   |   0    |

    Args:
        lookup_res: 字典树查表的结果。
        tag_dict: token tag 的所有映射表，使用 BIO，因为 BI可以解决字典树检索中重叠问题。
        sent_lens: 当前句长。要注意，一定要每个字符拆分开，因为字典树是存放的一个个 token，没有任何 span 信息。
        other_marker: 是否把找不到的所有 token 都标记为 Other tag。 True 时相当于第一种情况，非字典内的全是黑名单，False 时第二种情况，可以自己单独维护一个黑名单。

    Return:
        List
    """
    cgr = OrderedDict((k, [0] * sent_lens) for k in tag_dict)
    if other_marker:
        cgr['O'] = [1] * sent_lens
    for lookup in lookup_res:
        dict_name = list(lookup.keys())[0]
        spans = list(lookup.values())[0]
        for span in spans:
            cgr['B_' + dict_name][span[0]] = 1
            cgr['I_' + dict_name][span[0] + 1: span[1] + 1] = [1] * (span[1] - span[0])
            cgr['O'][span[0]: span[1] + 1] = [0] * (span[1] - span[0] + 1)
    return cgr


def main(input_file='test_data.txt', output_file='output_data.txt'):
    tokenizer = AutoTokenizer.from_pretrained('hfl/rbt3')
    tries = build_tries()
    tag_dict = codecs.open('tag.txt', 'r', 'utf-8').read().splitlines()
    output_file = codecs.open(output_file, 'w', 'utf-8')

    cnt = 0
    with codecs.open(input_file, 'r', 'utf-8') as read_file:
        for i, l in enumerate(read_file, 1):
            output = json.loads(l)
            cur_tokens = tokenizer.convert_ids_to_tokens(output['ids'])
            lookup_res = tries_search(tries, cur_tokens)
            cgr = convert_lookup_to_cgr(lookup_res, tag_dict, len(cur_tokens), other_marker=False)
            output['arcs'] = list(cgr.values())
            output_file.writelines(json.dumps(output, ensure_ascii=False))
            output_file.writelines('\n')
            cnt += 1
            print(cnt)

    output_file.close()


if __name__ == '__main__':
    main()
